from typing import override
from mcp.types import EmptyResult
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.client.sse import sse_client

from ..common import MCPServerConfig, ModelParameter 
from .basic import BaseMCPClient
from scorpio.core.common.logger import get_logger
logger = get_logger(__name__)

class DefaultClient(BaseMCPClient):
    """Implementation for a single MCP server."""
    def __init__(self, 
                mcp_server: str, 
                mcp_config: MCPServerConfig, 
                model_param: ModelParameter):
        super().__init__(mcp_server, mcp_config, model_param)
        self.tools = []
        self.prompts = []
    
    @override
    async def start(self)->bool:
        """
        Start to connect SSE MCP server
        """
        try:
            # Initialize MCP Server session
            if self.mcp_config.transportType == "stdio":
                self.session= await self.__stdio_start()
            elif self.mcp_config.transportType == "sse":
                self.session=await self.__sse_start()
            else:
                logger.error(f"server {self.server}: unsupported transport type {self.mcp_config.transportType}")
                return False
            """
            elif self.mcp_config.transportType == "streamable-http":
                # streamable-http support in fastmcp, not mcp official feature at present
                self.client=self.__streamable_start()
            """
            await self.session.initialize()
            self.tools= await self.session.list_tools()
            self.prompts= await self.session.list_prompts()
            logger.debug(f"Server {self.server}: Connect to MCP server with protocol {self.mcp_config.transportType} successfully.")
            return True
        except Exception as e:
            logger.error(f"Server {self.server}: SSE connection error: {str(e)}")
            return False
    
    async def __stdio_start(self)->ClientSession:
        """
        Start to connect STDIO MCP server
        """
        # to-do: need to validate the method 
        server_params = StdioServerParameters(
            command=self.mcp_config.command,
            args=self.mcp_config.args,
            env=self.mcp_config.env
        )
        stdio_read, stdio_write = await self.exit_stack.enter_async_context(stdio_client(server_params))
        return await self.exit_stack.enter_async_context(
            ClientSession(stdio_read, stdio_write))

    async def __sse_start(self)->ClientSession:
        """
        Start to connect SSE MCP server
        """
        headers=None
        auth_token=None
        env=self.mcp_config.env
        if env is not None:
            auth_token=env.get("authorization_key")
        if auth_token is not None:
            headers={
                "Authorization": auth_token,
            }
        sse_read, sse_write = await self.exit_stack.enter_async_context(
            sse_client(
                url=self.mcp_config.url,
                headers=headers,
            )
        )
        return await self.exit_stack.enter_async_context(ClientSession(sse_read, sse_write))
    
    @override
    async def ping(self):
        result=await self.session.send_ping()
        return isinstance(result, EmptyResult)


    def get_system_message(self):
        # Get system message - either from systemMessageFile, systemMessage, or default
        system_msg = "You are a assistant with ability to use tool with function call."
        for p in self.prompts:
            if p.name=="get_system_topic":
                system_msg=p.description
        messages=[]
        chosen_model=self.models_cfg
        if "systemMessageFile" in chosen_model:
            try:
                with open(chosen_model["systemMessageFile"], "r", encoding="utf-8") as f:
                    system_msg = f.read()
            except Exception as e:
                logger.warning(f"Failed to read system message file: {e}")
                # Fall back to direct systemMessage if available
                messages.append({"role": "system", "content": chosen_model.get("systemMessage", system_msg)})
        else:
            messages.append({"role": "system", "content": system_msg})
        if "systemMessageFiles" in chosen_model:
            for file in chosen_model["systemMessageFiles"]:
                try:
                    with open(file, "r", encoding="utf-8") as f:
                        system_msg = f.read()
                        messages.append({"role": "system", "content": "File: " + file + "\n" + system_msg})
                except Exception as e:
                    logger.warning(f"Failed to read system message file: {e}")
        return messages

    @override
    async def cleanup(self):
        """Clean up resources"""
        await self.exit_stack.aclose()
    """
    async def chat_loop(self):
        # Run an interactive chat loop

        print("\nMCP Client Started!")
        print("Type your queries or 'quit' to exit.")
        
        while True:
            try:
                query = input("\nQuery: ").strip()
                
                if query.lower() == 'quit':
                    break
                    
                response = await self.process_query(query)
                print("\n" + response)
                    
            except Exception as e:
                print(f"\nError: {str(e)}")
    """    

